# aem.two_d.single.constant_aquifer
#
# Copyright (c) 2000,2001 V.Kelson and M.Bakker

import copy
import math

import geometry
SearchableList = geometry.two_d.collections.SearchableList
XY = geometry.XY

import aem.container
ColumnList=aem.container.ColumnList
RowList=aem.container.RowList
TraceLine=aem.container.TraceLine

import aem.matrix

import base_aquifer

class ConstantAquifer(base_aquifer.Aquifer):
	"""Aquifer class with constant properties, a reference point and uniform flow

	Python call: ConstantAquifer(b = 0.0,H = 1.0,k = 1.0,n = 0.2,
		     refPt = geometry.two_d.Point(0,0),refHead = 1.0,
		     uniformFlow = XY(0,0),parent=None)

	Attributes:

	  - b: base elevation of aquifer
	  - H: thickness of aquifer
	  - k: hydraulic conductivity of aquifer
	  - n: porosity of aquifer
	  - refPt: location of reference point, XY object
	  - refHead: head at reference point
	  - uniformFlow: Qx,Qy components of uniform flow, XY object
	  - topElevation: elevation of top of aquifer
	  - topPotential: potential corresponding to head at top of aquifer
	  - parent: parent of the element, set to None (inherited from Element)
	  - parameters: ColumnList of strength parameters, to list of three entries,
	    the constant, and the components of uniform flow Qx and Qy
	  - solver: solver (inherited from Element)
	  - elementList: list of elements, contains all elements in aquifer (inherited from Element)

	Overloaded functions from Element: __init__, potentialInfluence, dischargeInfluence,
	getMatrixCoefficients, getMatrixContribution, check

	Overloaded from Aquifer: headToPotential, potentialToHead, transmissivity
	"""
	def __init__(self,b = 0.0,H = 1.0,k = 1.0,n = 0.2, \
		     refPt = geometry.two_d.Point(0,0),refHead = 1.0, \
		     uniformFlow = XY(0,0),parent=None):
		'''Constructor
		Accepts the aquifer properties and reference, reference point and uniform flow
		'''

		base_aquifer.Aquifer.__init__(self)
		self.b = b
		self.H = H
		self.k = k
		self.n = n
		self.refPt = copy.copy(refPt)
		self.refHead = refHead
		self.uniformFlow = copy.copy(uniformFlow)
		# Computed parameters
		self.topElevation = self.b + self.H
		self.topPotential = 0.5 * self.k * self.H ** 2
		# Strength parameters
		# Index  0 - Reference Constant
		#        1 - Qx
		#        2 - Qy
		#self.parameters = ColumnList([ self.headToPotential(self.refPt,self.refHead),
		#			       self.uniformFlow.x,
		#			       self.uniformFlow.y ] )
		self.parameters = ColumnList([ self.headToPotential(self.refPt,self.refHead),
					       self.uniformFlow.x,
					       self.uniformFlow.y ] )

	def __repr__(self):
		return 'ConstantAquifer' + str((self.b,self.H,self.k,self.n,self.refPt,self.refHead,self.uniformFlow))

	def potentialInfluence(self,xy,t=0):
		dxy = xy - self.refPt
		return RowList( [ 1.0, -dxy.x, -dxy.y ])

	def dischargeInfluence(self,xy,t=0):
		return RowList( [ XY(0.0,0.0), XY(1.0,0.0), XY(0.0,1.0) ] )

	def extractionInfluence(self,t=0):
		return RowList( [ 0.0, 0.0, 0.0 ] )

	def rechargeInfluence(self,pt,t=0):
		return RowList( [ 0.0, 0.0, 0.0 ] )

	def getMatrixCoefficients(self,fa,ctlPt):
		# Returns a single matrix coefficent based on the reference head, at the point ctlPt
		return [ fa(self,ctlPt)[0] ]

	def getMatrixContribution(self,elementList):
		# Returns the MatrixRow objects for this element
		# This element solves only for the reference point constant (influence element 0)
		mr = aem.matrix.MatrixRow()
		mc = aem.matrix.MatrixColumn(self,0)
		pt = self.refPt
		for e in elementList:
			mr.A = mr.A + e.getMatrixCoefficients(lambda z,p:z.potentialInfluence(p),pt)
		mr.b = self.headToPotential(self.refPt,self.refHead) - \
		       self.potential(self.refPt)
		return ( [mr],[mc] )

	def check(self):
		return [ self.head(self.refPt) - self.refHead ]


	def headToPotential(self,pt,head):
		#Convert head to potential at the point pt
		if head > self.topElevation:
			# Confined flow
			rv = self.k * self.H * (head-self.b) - self.topPotential
		else:
			rv = 0.5 * self.k * (head-self.b) ** 2
		return rv

	def potentialToHead(self,pt,pot):
		#Convert head to potential at the point pt
		if pot >= self.topPotential:
			rv = self.b + (pot + self.topPotential) / (self.k * self.H)
		else:
			if pot <= 0:
				rv = self.b
			else:
				rv = self.b + math.sqrt( 2.0 * pot / self.k )
		return rv

	def transmissivity(self,pt,t=0):
		#transmissivity at the point pt
		head=self.head(pt,t)
		if head > self.topElevation:
			rv = self.k * self.H
		elif head > self.b:
			rv = self.k * (head-self.b)
		else:
			rv = 0
		return rv

	def velocity(self,xyzpt,t=0):
		pt = xyzpt.XY()
		QxQy = self.discharge(pt)
		h = self.head(pt)
		if h < self.topElevation:
			qxqy = QxQy / (h - self.b)
			qz = -( self.recharge(pt) + \
				( QxQy.x ** 2 + QxQy.y ** 2 )/( self.k * h **2 ) ) * \
				( xyzpt.z - self.b ) / h
		else:
			qxqy = QxQy / self.H
			qz = - self.recharge(pt) * (xyzpt.z - self.b) / self.H
		rv = geometry.XYZ(qxqy.x,qxqy.y,qz) / self.n
		return rv

	def trace(self,xyzstart,step,tmax,maxsteps=10):
		'''Tracing routine

		Input:
		- xyzstart: XYZ object
		- step: 3D step size
		- tmax: maximum time (will stop at tmax if no other element has been reached)
		- maxsteps: maximum number of steps (default = 10)

		Returns:
		- TraceLine object'''

		# Check starting point
		h = self.head(xyzstart.XY())
		if xyzstart.z < self.b or xyzstart.z > min(h,self.topElevation):
			raise 'Starting point outside saturated aquifer'

		traceLine = TraceLine()
		traceLine = traceLine.addpoint(xyzstart,0)

		xyzold = xyzstart
		told = 0
		nstep = 0

		while nstep < maxsteps:
			# Compute velocity and stepsize
			vxyz = self.velocity(xyzold)
			velo = vxyz.magnitude()
			tstep = step / velo
			# Compute new point
			xyznew = xyzold + vxyz * tstep
			tnew = told + tstep
			changed = 0
			# Check whether point is at or near element
			for e in self.elementList:
				check = e.nearElement(xyzold,xyznew,step)
				if check.change:
					xyznew = check.xyzchanged
					stepnew = (xyznew-xyzold).magnitude()
					tnew = told + stepnew/step * tstep + check.delt
					changed = 1
				if check.stop: break
			# Check whether point is inside aquifer
			# Does step out somewhat crudely, since it may step above water table
			if self.recharge(xyzold.XY()) < 0:
				if xyznew.z > min( self.topElevation, self.head(xyznew.XY()) ):
					check.stop = 1
			# If no change, do corrector step
			if not changed and not check.stop:
				vxyznew = self.velocity(xyznew)
				vxyz = 0.5 * (vxyz+vxyznew)
				velo = vxyz.magnitude()
				tstep = step / velo
				xyznew = xyzold + vxyz * tstep
				tnew = told + tstep
				# Check again whether is at or near point
				for e in self.elementList:
					check = e.nearElement(xyzold,xyznew,step)
					if check.change:
						xyznew = check.xyzchanged
						stepnew = (xyznew-xyzold).magnitude()
						tnew = told + stepnew/step * tstep + check.delt
					if check.stop: break
			# Check whether point is inside aquifer
			# Does step out somewhat crudely, since it may step above water table
			if self.recharge(xyzold.XY()) < 0:
				if xyznew.z > min( self.topElevation, self.head(xyznew.XY()) ):
					check.stop = 1
			# Check whether maximum time has been reached
			if tnew > tmax:
				xyznew = xyzold + (tmax-told) * vxyz
				tnew = tmax
				check.stop = 1
			# Add point to trace line
			traceLine = traceLine.addpoint(xyznew,tnew)
			nstep = nstep + 1
			if check.stop: break
			xyzold=xyznew; told=tnew

		if tnew == tmax:
			print "Reached maximum time"
		if nstep >= maxsteps:
			print "Reached maximum number of steps"
		if check.stop and check.change:
			print "Reached element"

		return traceLine

	def tracewrite(self,trace):
		'''Routine to write trace (TraceLine instance) to file'''
		out = open('/temp/trace.bln','w')
		out.write(str(len(trace.xyzt))+' 1\n')
		for i in range(len(trace.xyzt)):
			out.write(str(trace.xyzt[i].x)+' '+str(trace.xyzt[i].y)+'\n')
		out.close()

	def tracelines(self,xyz,step,tmax,Nmax,filename):
		'''Routine for calculating multiple tracelines and writing them to file.
		xyz is list of XYZ objects, filename is character string, such as 'mark.dat'
		and must be between quotes
		'''
		out = open(filename,'w')
		for i in range(len(xyz)):
			trace=self.trace(xyz[i],step,tmax,Nmax)
			out.write(str(len(trace.xyzt))+' 1\n')
			for i in range(len(trace.xyzt)):
				out.write(str(trace.xyzt[i].x)+' '+str(trace.xyzt[i].y)+'\n')
		out.close()

	def makegrid(self,xmin,xmax,xstep,ymin,ymax,ystep,filename,t=0):
		'''Routine for calculating grid of heads and write surfer grid
		filename must be character string, such as /temp/tim.grd (forward slashes!)'''
		nx=int((xmax-xmin)/xstep)
		ny=int((ymax-ymin)/ystep)
		matrix=[]
		zmin = self.head(XY(xmin,ymin)); zmax = zmin
		for j in range(ny+1):
			row = []
			y=ymin + j*ystep
			for i in range(nx+1):
				x=xmin + i*xstep;
				h = self.head(XY(x,y),t)
				if h < zmin: zmin = h
				if h > zmax: zmax = h
				row = row + [ h ]
			matrix = matrix + [row]

		out = open(filename,'w')
		out.write('DSAA\n')

		out.write(str(nx+1)+' '+str(ny+1)+'\n')
		out.write(str(xmin)+' '+str(xmin+nx*xstep)+'\n')
		out.write(str(ymin)+' '+str(ymin+ny*ystep)+'\n')
		out.write(str(zmin)+' '+str(zmax)+'\n')

		for j in range(ny+1):
			row = matrix[j]
			for i in range(nx):
				out.write(str(row[i])+' ')
			out.write(str(row[nx])+'\n')
		out.close()
